# -*- coding:utf-8 -*-

import logging
from datetime import datetime

import pytz

from plugin.configs import configs
from plugin.db.mongo_base import MongoBase

tz_us = pytz.timezone('America/New_York')
tz_cn = pytz.timezone('Asia/Shanghai')


def mongo_time(dtime, timezone=tz_cn):
    return timezone.localize(dtime)


class MongoCollection(object):
    def __init__(self, uri=None, db=None):
        self.logger = logging.getLogger('MongoCollection')
        self.uri = uri
        self.db = db

    def mongo_conn(self, uri=None, db=None):
        if self.uri and self.db:
            uri = self.uri
            db = self.db
        if uri and db:
            try:
                mongo_connection = MongoBase(uri, database=db, appname='pyspider_downloader_{}'.format(db)).db

            except Exception as e:
                self.logger.error(e)
                return None
        else:
            try:
                uri = configs['mongo_or']['mongo_uri']
                db = configs['mongo_or']['mongo_db']
                mongo_connection = MongoBase(uri, database=db, appname='pyspider_downloader_{}'.format(db)).db
            except Exception as e:
                self.logger.error(e)
                return None
        return mongo_connection


class MongoDBTool(object):
    def __init__(self, uri=None, db=None):
        self.logger = logging.getLogger('mongoDBTool')
        self.mongo_conn = MongoCollection(uri, db).mongo_conn()

    def save_data(self, data, table_name):
        getattr(self.mongo_conn, table_name).insert(data)

    def delete_data(self, find, table_name):
        getattr(self.mongo_conn, table_name).delete_one(find)

    def save_mongodb(self, data, find, tablename):
        """
        检查key值是否存在，存在则更新，不存在则插入
        :param data: 需要插入的数据dict
        :param find: 匹配的字段dict
        :param tablename: 相关表
        :param select_field: 选择特定字段
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            getattr(self.mongo_conn, tablename).update(find, {'$set': data}, multi=True)
        else:
            getattr(self.mongo_conn, tablename).insert(data)

    def save_mongodb_time(self, data, find, tablename, convert_tzone=True):
        """
        检查key值是否存在，存在则更新，不存在则插入, 在save_mongodb基础上加入create_time,和update_time
        :param data: 需要插入的数据dict
        :param find: 匹配的字段dict
        :param tablename: 相关表
        :param is_china_timezone: 是否中国区域
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        ts = datetime.now()
        print(convert_tzone, ts)
        # 中国区，转化utf-8时间
        if convert_tzone:
            # 行研
            self.transform_mongotime(data, 'time')
            self.transform_mongotime(data, 'write_time')
            # EDB
            # self.transform_mongotime(data, 'format_time')
            # self.transform_mongotime(data, 'published_time')

            self.transform_mongotime(data, 'downloaded_time')
            data['update_time'] = mongo_time(ts)
            data['create_time'] = mongo_time(ts)
        else:
            data['create_time'] = ts
            data['update_time'] = ts

        if result.count() >= 1:
            # 更新操作，排除create_time
            for pop_key in ['create_time']:
                if pop_key in data.keys():
                    data.pop(pop_key)
            getattr(self.mongo_conn, tablename).update(find, {'$set': data}, multi=True)
        else:
            getattr(self.mongo_conn, tablename).update(find, {'$set': data}, upsert=True)
        return True

    @staticmethod
    def transform_mongotime(data, key):
        v = data.get(key)
        if v is None or v == '':
            return
        try:
            data[key] = mongo_time(v)
        except Exception as ex:
            pass

    def save_mongodb_thomson_fix(self, data, find, tablename):
        """
        专用用于thomson下载pdf_list
        检查key值是否存在，存在则更新，不存在则插入
        :param data: 需要插入的数据dict
        :param find: 匹配的字段dict
        :param tablename: 相关表
        :param select_field: 选择特定字段
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            data = {k:v for k,v in list(data.items()) if k != 'is_uploader_aliyun'}
            getattr(self.mongo_conn, tablename).update(find, {'$set':data}, multi=True)
        else:
            getattr(self.mongo_conn, tablename).insert(data)

    def update_mongodb(self, data, find, tablename, multi=False):
        """
        检查key值是否存在，存在则更新，不存在不插入
        :param data:
        :param find:
        :param tablename:
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            getattr(self.mongo_conn, tablename).update(find, {'$set': data}, multi=multi)

    def select(self, tablename, find):
        return getattr(self.mongo_conn, tablename).find(find)

    def select_field(self, tablename, find, select_field):
        return getattr(self.mongo_conn, tablename).find(find, select_field)

    def select_field_limit(self, tablename, find, select_field, limit):
        return getattr(self.mongo_conn, tablename).find(find, select_field).limit(limit)

    def save_mongodb_field(self, data, find, tablename, field, option):
        """
        通过find查看数据库是否存在对应数据，若存在则更新指定字段
        :param data: 需要插入的数据
        :param find: 查找存在数据的条件
        :param tablename: 更新表
        :param field: 指定字段,类型为list或者tuple
        :param option: boolean类型，表示是更新相关字段还是不更新相关字段,True表示更新field相关字段
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            if option:
                getattr(self.mongo_conn, tablename).update(find,
                                                           {'$set': {k: v for k, v in list(data.items()) if k in field}}, multi=True)
            else:
                getattr(self.mongo_conn, tablename).update(find,
                                                           {'$set': {k: v for k, v in list(data.items()) if k not in field}}, multi=True)
        else:
            getattr(self.mongo_conn, tablename).insert(data)

    def save_mongodb_field_fix(self, data, find, tablename, field, option):
        """
        修正了在做杭研数据中create_time和update_time与实际插入时间不匹配的问题
        修正之后：update_time为插入数据库的实际时间
        1.数据库中存在指定记录，create_time不会被修改，而update_time会更新到指定时间
        2.数据库中不存在指定记录，create_time和update_time都为当前时间，且一致

        通过find查看数据库是否存在对应数据，若存在则更新指定字段
        :param data: 需要插入的数据
        :param find: 查找存在数据的条件
        :param tablename: 更新表
        :param field: 指定字段,类型为list或者tuple
        :param option: boolean类型，表示是更新相关字段还是不更新相关字段,True表示更新field相关字段
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            data['update_time'] = mongo_time(datetime.now())
            if option:
                getattr(self.mongo_conn, tablename).update(find,
                                                           {'$set': {k: v for k, v in list(data.items()) if k in field}}, multi=True)
            else:
                getattr(self.mongo_conn, tablename).update(find,
                                                           {'$set': {k: v for k, v in list(data.items()) if k not in field}}, multi=True)
        else:
            data['create_time'] = mongo_time(datetime.now())
            data['update_time'] = mongo_time(datetime.now())

            getattr(self.mongo_conn, tablename).insert(data)

    def count_table_num(self, tablename, find):
        """
        通过给到筛选条件，统计表中的数量
        :param tablename:
        :param find:
        :return:
        """
        return getattr(self.mongo_conn, tablename).find(find).count()

    def save_item_time(self, data, find, tablename):
        """
        检查key值是否存在，存在跳过，不存在则插入, 在save_mongodb基础上加入create_time,和update_time
        :param data: 需要插入的数据dict
        :param find: 匹配的字段dict
        :param tablename: 相关表
        :param select_field: 选择特定字段
        :return:
        """
        result = getattr(self.mongo_conn, tablename).find(find)
        if result.count():
            pass
        else:
            data['create_time'] = mongo_time(datetime.now())
            data['update_time'] = mongo_time(datetime.now())
            getattr(self.mongo_conn, tablename).insert(data)

    def remove(self, find, tablename):
        getattr(self.mongo_conn, tablename).remove(find)




